import cv2
import numpy
import win32ui
import win32api
import win32con
import win32gui
import ctypes
import tkinter as tk

class AreaSelector:
    def __init__(self):
        # 设置DPI感知处理系统缩放问题（重要！）
        ctypes.windll.user32.SetProcessDPIAware()

    def on_mouse_press(self, event):
        self.start_x = event.x
        self.start_y = event.y

    def on_mouse_move(self, event):
        self.end_x = event.x
        self.end_y = event.y
        # 更新矩形选区
        self.canvas.delete("rect")
        self.canvas.create_rectangle(
            self.start_x, self.start_y, self.end_x, self.end_y,
            outline="#ff0000", width=2, tags="rect"
        )

    def on_mouse_release(self, event):
        self.root.destroy()

    def screen_capture(self):
        """ @return: matrix """
        # 获取桌面
        hdesktop = win32gui.GetDesktopWindow()
        # 分辨率适应
        width = win32api.GetSystemMetrics(win32con.SM_CXVIRTUALSCREEN)
        height = win32api.GetSystemMetrics(win32con.SM_CYVIRTUALSCREEN)
        left = win32api.GetSystemMetrics(win32con.SM_XVIRTUALSCREEN)
        top = win32api.GetSystemMetrics(win32con.SM_YVIRTUALSCREEN)
        # 创建设备描述表
        desktop_dc = win32gui.GetWindowDC(hdesktop)
        img_dc = win32ui.CreateDCFromHandle(desktop_dc)
        mem_dc = img_dc.CreateCompatibleDC()
        # 创建位图对象
        bitmap = win32ui.CreateBitmap()
        bitmap.CreateCompatibleBitmap(img_dc, width, height)
        mem_dc.SelectObject(bitmap)
        mem_dc.BitBlt((0, 0), (width, height), img_dc, (0, 0), win32con.SRCCOPY)
        bmpstr = bitmap.GetBitmapBits(True)

        img_cv = numpy.frombuffer(bmpstr, dtype = 'uint8')
        img_cv.shape = (height, width, 4)
        img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGRA2BGR)

        # 内存释放
        mem_dc.DeleteDC()
        win32gui.DeleteObject(bitmap.GetHandle())

        return img_cv

    def wait_select(self):
        self.root = tk.Tk()
        self.root.attributes('-fullscreen', True)  # 设置为全屏
        self.root.attributes('-alpha', 0.3)  # 设置透明度以便可以看到下面的内容

        # 创建一个画布，覆盖整个窗口
        self.canvas = tk.Canvas(self.root, cursor="cross")
        self.canvas.pack(fill=tk.BOTH, expand=True)

        # 绑定鼠标事件
        self.canvas.bind("<ButtonPress-1>", self.on_mouse_press)
        self.canvas.bind("<B1-Motion>", self.on_mouse_move)
        self.canvas.bind("<ButtonRelease-1>", self.on_mouse_release)

        # 显示一个提示信息
        label = tk.Label(self.root, text="请拖拽鼠标选择截图区域", font=('Arial', 16))
        label.pack()

        # loop等待销毁
        self.root.mainloop()

        left = min(self.start_x, self.end_x)
        top = min(self.start_y, self.end_y)
        right = max(self.start_x, self.end_x)
        bottom = max(self.start_y, self.end_y)

        screenshot = self.screen_capture()
        crop_pic = screenshot[top:bottom, left:right]

        return (left, top, right, bottom), crop_pic

